package com.ruoyu.cloud.ai.controller;

import com.ruoyu.cloud.ai.splitter.QjcTextSplitter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.InputStreamResource;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.ai.document.Document;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@RestController
@Slf4j
public class DocumentController {


    private final EmbeddingModel embeddingModel;
    private final VectorStore vectorStore;

    @Autowired
    public DocumentController(EmbeddingModel embeddingModel, VectorStore vectorStore) {
        this.embeddingModel = embeddingModel;
        this.vectorStore = vectorStore;
    }

    @GetMapping("/ai/embedding")
    public Map embed(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        EmbeddingResponse embeddingResponse = this.embeddingModel.embedForResponse(List.of(message));
        return Map.of("embedding", embeddingResponse);
    }

    //localhost:8081/embedding
    @GetMapping("/embedding")
    public Object embeddingTest() {
        // 文本嵌入
        float[] embed = embeddingModel.embed("你好，我是小码农叔叔");
        log.info("文本转换得到的向量: {}", embed);
        return embed;
    }

    @PostMapping("read/multipart-file")
    public List<Document> readMultipartFile(@RequestParam MultipartFile file) {
        // 从IO流中读取文件
        try {
            Resource resource = new InputStreamResource(file.getInputStream());
            List<Document> documents = new TikaDocumentReader(resource).read();
            return documents;
        } catch (Exception e) {
        }
        return null;

    }

    @Value("classpath:meituan-qa.txt")
    private Resource resource;

    /*** 从本地文件读取* @return*/
    @GetMapping("/read/from-local-file")
    public Object readLocalFile() {
//        FileSystemResource resource = new FileSystemResource("classpath:meituan-qa.txt");
        List<Document> documents = new TikaDocumentReader(resource).read();
        return documents;
    }

    /*** 分割后的文档列表  localhost:8081/doc/split* @param file* @return* @throws Exception*/
    @PostMapping("/doc/split")
    public Object docSplit(@RequestParam MultipartFile file) throws Exception {
        InputStreamResource inputStreamResource = new InputStreamResource(file.getInputStream());
        List<Document> documentList = new TikaDocumentReader(inputStreamResource).read();
        // 分割后的文档列表
        List<Document> splitDocList = new TokenTextSplitter().split(documentList);
        List<String> contentList = splitDocList.stream().map(Document::getText).collect(Collectors.toList());
        return contentList;
    }

    /*** 分割后的文档列表  localhost:8081/doc/split* @param file* @return* @throws Exception*/
    @GetMapping("/local/split")
    public Object localSplit() throws Exception {
        List<Document> documentList = new TikaDocumentReader(resource).read();
        // 分割后的文档列表
        QjcTextSplitter qjcTextSplitter = new QjcTextSplitter();
        List<Document> splitDocList = qjcTextSplitter.apply(documentList);
//        List<Document> splitDocList = new TokenTextSplitter().split(documentList);
        List<String> contentList = splitDocList.stream().map(Document::getText).collect(Collectors.toList());
        return splitDocList;
    }

    /*** 文档保存到向量数据库   localhost:8081/doc/embedding* @param file* @return* @throws Exception*/
    @PostMapping("/doc/embedding")
    public Object docEmbedding(@RequestParam MultipartFile file) throws Exception {
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(new InputStreamResource(file.getInputStream()));
        List<Document> splitDocuments = new TokenTextSplitter().apply(tikaDocumentReader.read());
        // 存入向量数据库，这个过程会自动调用embeddingModel,将文本变成向量存入
        vectorStore.add(splitDocuments);
        return true;
    }

    /*** 文档保存到向量数据库   localhost:9091/doc/local/embedding* @param file* @return* @throws Exception*/
    @GetMapping("/doc/local/embedding")
    public Object docLocalEmbedding() throws Exception {
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(resource);
        List<Document> documentList = tikaDocumentReader.read();
//        List<Document> splitDocuments = new TokenTextSplitter().apply(tikaDocumentReader.read());
        // 存入向量数据库，这个过程会自动调用embeddingModel,将文本变成向量存入
        QjcTextSplitter qjcTextSplitter = new QjcTextSplitter();
        List<Document> splitDocuments = qjcTextSplitter.apply(documentList);
        vectorStore.add(splitDocuments);
        return splitDocuments;
    }


    @GetMapping("/search")
    public Object search(String msg) throws Exception {
//        SearchRequest searchRequest = SearchRequest.builder()
//                .query(msg)
//                .topK(10).build();
        List<Document> documents = vectorStore.similaritySearch(msg);
        return documents;
    }

}
